from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.exceptions import RequestValidationError
from pydantic import BaseModel, validator
import requests
import os
from dotenv import load_dotenv
from lib.mysql import MySQLConnection
import json
from fish_audio_sdk import Session, TTSRequest
import io
import time
from tos import TosClientV2

# 加载环境变量
load_dotenv()

app = FastAPI()


db = MySQLConnection(host='mysql-116593bb4c5e-public.rds.volces.com', database='uaccount', user='omelette', password='qwer@123123')
connection = db.get_connection()

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class VoiceCloneRequest(BaseModel):
    audio_url: str
    text: str = ""  # 可选参数，默认为空字符串
    title: str = "Demo"  # 可选参数，默认为"Demo"

    @validator('audio_url')
    def validate_audio_url(cls, v):
        if not v:
            raise ValueError('音频URL不能为空')
        if not v.startswith(('http://', 'https://')):
            raise ValueError('无效的音频URL格式')
        return v

# 添加局异常处理器
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    print(f"Validation error: {exc.errors()}")
    return JSONResponse(
        status_code=200,
        content={
            "data": {
                "code": 400,
                "message": "请求参数错误: " + str(exc.errors()[0].get('msg', '未知错误')),
                "data": None
            }
        }
    )

async def download_audio(url: str) -> bytes:
    try:
        response = requests.get(url)
        response.raise_for_status()
        return response.content
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"音频文件下载失败: {str(e)}")

@app.post("/voice-clone")
async def voice_clone(request: VoiceCloneRequest):
    try:
        print(f"Received request - URL: {request.audio_url}, Text: {request.text}")
        
        # 下载音频文件
        audio_content = await download_audio(request.audio_url)
        # 准备请求数据
        files = [
            ("voices", ("voice.mp3", audio_content, "audio/mpeg"))
        ]
        
        data = [
            ("visibility", "private"),
            ("type", "tts"),
            ("title", request.title),
            ("train_mode", "fast"),
            ("enhance_audio_quality", "true")
        ]

        # 如果提供了文本，添加到请求数据中
        if request.text:
            data.append(("texts", request.text))

        print(os.getenv('AUDIO_FISH_API_KEY'))
        # 发送请求到 Fish Audio API
        response = requests.post(
            "https://api.fish.audio/model",
            files=files,
            data=data,
            headers={
                "Authorization": f"Bearer {os.getenv('AUDIO_FISH_API_KEY')}"
            }
        )

        response_data = response.json()
        print(f"Fish Audio API response: {response_data}")

        # 插入数据库
        cursor = connection.cursor()
        sql = """
            INSERT INTO fish_voice_data (userid, reference_id, title, audio_url, response_data)
            VALUES (%s, %s, %s, %s, %s)
        """
        # 这里的 userid 需要从请求中获取，暂时使用默认值
        userid = "default_user"  # 你需要根据实际情况获取用户ID
        cursor.execute(sql, (
            userid, 
            response_data['_id'],
            response_data['title'],  # 从���应数据中获取 title
            request.audio_url,
            json.dumps(response_data)  # 将完整响应数据转换为JSON字符串
        ))
        connection.commit()
        cursor.close()

        # 返回结果
        return {
            "data": {
                "code": 0,
                "message": "语音克隆请求成功",
                "data": response_data
            }
        }

    except Exception as e:
        if 'cursor' in locals():
            cursor.close()
        if connection.is_connected():
            connection.rollback()
        print(f"Error in voice clone: {str(e)}")
        return JSONResponse(
            status_code=200,
            content={
                "data": {
                    "code": 500,
                    "message": f"语音克隆失败: {str(e)}",
                    "data": None
                }
            }
        )

# 添加新的请求模型
class TTSRequestModel(BaseModel):
    reference_id: str
    text: str

    @validator('reference_id')
    def validate_reference_id(cls, v):
        if not v:
            raise ValueError('reference_id不能为空')
        return v

    @validator('text')
    def validate_text(cls, v):
        if not v:
            raise ValueError('合成文本不能为空')
        return v

# 初始化火山存储客户端
bucket_name = os.getenv('TOS_BUCKET_NAME')
endpoint = os.getenv('TOS_ENDPOINT')
TosClient = TosClientV2(
    ak=os.getenv('TOS_ACCESS_KEY'),
    sk=os.getenv('TOS_SECRET_KEY'),
    endpoint=endpoint,
    region=os.getenv('TOS_REGION')
)

@app.post("/tts")
async def text_to_speech(request: TTSRequestModel):
    start_time = time.time()
    try:
        print(f"[TTS] 开始处理请求 - reference_id: {request.reference_id}, text: {request.text}")
        
        # 初始化 Fish Audio session
        session = Session(os.getenv('AUDIO_FISH_API_KEY'))
        tts_start_time = time.time()
        print(f"[TTS] 初始化 Session 耗时: {(tts_start_time - start_time):.2f}秒")
        
        # 创建内存缓冲区来存储音频数据
        audio_buffer = io.BytesIO()
        
        # 执行 TTS 请求
        print("[TTS] 开始生成音频...")
        for chunk in session.tts(TTSRequest(
            reference_id=request.reference_id,
            text=request.text
        )):
            audio_buffer.write(chunk)
        
        tts_end_time = time.time()
        print(f"[TTS] 音频生成完成，耗时: {(tts_end_time - tts_start_time):.2f}秒")
        
        # 将指针移到开始位置
        audio_buffer.seek(0)
        
        # 生成唯一的文件名
        filename = f"tts/{request.reference_id}/{int(time.time())}.mp3"
        
        # 上传到火山存储
        print("[TTS] 开始上传音频文件...")
        upload_start_time = time.time()
        
        response = TosClient.put_object(
            bucket=bucket_name,
            key=filename,
            content=audio_buffer.getvalue()
        )
        
        upload_end_time = time.time()
        print(f"[TTS] 文件上传完成，耗时: {(upload_end_time - upload_start_time):.2f}秒")
        
        if response.status_code == 200:
            # 生成可访问的URL
            url = f"https://{bucket_name}.{endpoint}/{filename}"
            total_time = time.time() - start_time
            print(f"[TTS] 请求处理完成，总耗时: {total_time:.2f}秒")
            
            return {
                "data": {
                    "code": 0,
                    "message": "TTS转换成功",
                    "data": {
                        "url": url,
                        "process_time": {
                            "tts_time": round(tts_end_time - tts_start_time, 2),
                            "upload_time": round(upload_end_time - upload_start_time, 2),
                            "total_time": round(total_time, 2)
                        }
                    }
                }
            }
        else:
            raise Exception("文件上传失败")

    except Exception as e:
        error_time = time.time() - start_time
        print(f"[TTS] 处理失败，耗时: {error_time:.2f}秒，错误: {str(e)}")
        return JSONResponse(
            status_code=200,
            content={
                "data": {
                    "code": 500,
                    "message": f"TTS转换失败: {str(e)}",
                    "data": None
                }
            }
        )

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=3000) 